import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.ensemble import RandomForestClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import classification_report
from joblib import dump
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score, classification_report

class Classification:
    def __init__(self):
        self.df = pd.read_csv('daily.csv')
        print("数据列名:", self.df.columns.tolist())  # 打印列名确认

    def get_condition(self):
        df = self.df
        
        # 使用正确的列名计算比率
        df['new_ratio'] = df['max_close'] / df['the_slose']
        df['min_ratio'] = df['min_close'] / df['the_slose']

        high_return_threshold = df['new_ratio'].quantile(0.4)
        high_risk_threshold = df['min_ratio'].quantile(0.4)

        conditions = [
            (df['new_ratio'] >= high_return_threshold) & (df['min_ratio'] <= high_risk_threshold),
            (df['new_ratio'] >= high_return_threshold) & (df['min_ratio'] > high_risk_threshold),
            (df['new_ratio'] < high_return_threshold) & (df['min_ratio'] > high_risk_threshold),
            (df['new_ratio'] < high_return_threshold) & (df['min_ratio'] < high_risk_threshold),
        ]
        
        labels = ['高收益高风险', '高收益低风险', '低收益低风险', '低收益高风险']
        df['category'] = np.select(conditions, labels, default='未知')

        features = df[['eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin',
                      'fcff', 'fcfe', 'bps', 'grossprofit_margin']]

        le = LabelEncoder()
        df['category_crossed'] = le.fit_transform(df['category'])
        
        X_train, X_test, y_train, y_test = train_test_split(
            features, df['category_crossed'], test_size=0.3, random_state=24)
            
        return X_train, X_test, y_train, y_test, le

    def knn_utils(self, X_train, X_test, y_train, y_test, le):
        knn = KNeighborsClassifier(n_neighbors=5)
        knn.fit(X_train, y_train)
        y_pred = knn.predict(X_test)
        
        dump(knn, 'knn_classifier.joblib')
        dump(le, 'label_encoder.joblib')
        
        print(classification_report(y_test, y_pred, target_names=le.classes_))

    def svc_utils(self, X_train, X_test, y_train, y_test):
        svc = SVC()
        svc.fit(X_train, y_train)
        predictions = svc.predict(X_test)
        print("Accuracy score: %.4f" % accuracy_score(predictions, y_test))
        print("Classification report for classifier:\n%s\n" % classification_report(y_test, predictions))

if __name__ == '__main__':
    ci = Classification()
    X_train, X_test, y_train, y_test, le = ci.get_condition()
    ci.knn_utils(X_train, X_test, y_train, y_test, le)